{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%sh\n",
    "pip install -qU langchain_community yfinance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pprint\n",
    "\n",
    "import yfinance\n",
    "from langchain_community.llms import Ollama"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "AGENT_MODEL_NAME = \"arcee-ai/arcee-agent\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = Ollama(\n",
    "    model=AGENT_MODEL_NAME,\n",
    "    system=\"\"\"You are a helpful chat assistant.\"\"\",\n",
    "    top_p=0.7,\n",
    ")  # assuming you have run `ollama pull arcee-ai/arcee-agent`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "func_prompt = \"\"\"You have four primary functions: checking the last price of a specified stock, finding the name of a company CEO, finding what a company does, and answering specific questions about a company. Use the appropriate function based on the user's query.\n",
    "\n",
    "### Functions:\n",
    "\n",
    "1. **get_stock_price(company_name: str, stock_symbol: str) -> str**\n",
    "    - This function returns the last close price of a specified stock.\n",
    "    - Input: A company name (e.g., Mc Donalds), which you must convert to a stock symbol (e.g., MCD).\n",
    "    - Output: A string containing the last close price of the specified stock (e.g., \"The last closing price of Mc Donalds (MCD) is $250.00\").\n",
    "\n",
    "2. **get_ceo_name(company_name: str, stock_symbol: str) -> str**\n",
    "- This function returns the name of the CEO of a specified company.\n",
    "- Input: A company name (e.g., Mc Donalds), which you must convert to a stock symbol (e.g., MCD).\n",
    "- Output: A string containing the name of the CEO of the specified company (e.g., \"The CEO of Mc Donalds is John Doe\").\n",
    "\n",
    "3. **get_company_summary(company_name: str, stock_symbol: str) -> str**\n",
    "- This function returns a summary describing the business activities of a specified company.\n",
    "- Input: A company name (e.g., Mc Donalds), which you must convert to a stock symbol (e.g., MCD).\n",
    "- Output: A string containing a detailed summary of the specified company's business activities.\n",
    "\n",
    "4. **answer_general_question(question: str) -> str**\n",
    "    - This function answers questions in general.\n",
    "    - Input: a user question.\n",
    "    - Output: your best answer to the user question.\n",
    "\n",
    "### Instructions:\n",
    "\n",
    "- If the user asks a question related to the price of a stock, use the `get_stock_price` function.\n",
    "- If the user asks a question related to the CEO of a company, use the `get_ceo_name` function.\n",
    "- If the user asks a general question about a company's activities, use the `get_company_summary` function.\n",
    "- If the user asks any other question, use the `answer_general_question` function. Only return the result of the function call, not your internal reasoning.\n",
    "- Only return the result of the function call.\n",
    "\n",
    "### User Query:\\n\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_stock_price(company_name, stock_symbol):\n",
    "    stock = yfinance.Ticker(stock_symbol)\n",
    "    price = stock.history(period=\"1d\")[\"Close\"].values[0]\n",
    "    return (\n",
    "        f\"The last closing price of {company_name} ({stock_symbol}) was ${price:.2f}.\"\n",
    "    )\n",
    "\n",
    "\n",
    "def get_ceo_name(company_name, stock_symbol):\n",
    "    stock = yfinance.Ticker(stock_symbol)\n",
    "    info = stock.info\n",
    "    ceo = info[\"companyOfficers\"][0][\"name\"]\n",
    "    return f\"The CEO of {company_name} is {ceo}. The full job title is {info['companyOfficers'][0]['title']}.\"\n",
    "\n",
    "\n",
    "def get_company_summary(company_name, stock_symbol):\n",
    "    stock = yfinance.Ticker(stock_symbol)\n",
    "    summary = stock.info[\"longBusinessSummary\"]\n",
    "    return (\n",
    "        f\"{company_name} ({stock_symbol}) is a company that is involved in {summary}.\"\n",
    "    )\n",
    "\n",
    "\n",
    "def answer_general_question(question):\n",
    "    ans = agent.invoke(f\"\"\"{question}\\n\"\"\")\n",
    "    return ans\n",
    "\n",
    "\n",
    "def function_call(input, model=agent, func_prompt=func_prompt):\n",
    "    \"\"\"\n",
    "    Interacts with an AI model by combining user input with a custom prompt.\n",
    "\n",
    "    Args:\n",
    "        input (str): The user's input string.\n",
    "        model (object, optional): The AI model to invoke. Defaults to agent.\n",
    "        func_prompt (str, optional): The custom prompt to combine with the user's input. Defaults to func_prompt.\n",
    "\n",
    "    Returns:\n",
    "        str: The response generated by the AI model in response to the combined prompt and input.\n",
    "    \"\"\"\n",
    "    payload = f\"\"\"{func_prompt}{input}\\n\"\"\"\n",
    "    ans = model.invoke(payload)\n",
    "    return ans\n",
    "\n",
    "\n",
    "def llm_pack(input):\n",
    "    \"\"\"\n",
    "    Evaluates a string generated by an AI model and returns the result.\n",
    "\n",
    "    This function wraps around 'function_call' to convert the AI's response into a executable code snippet,\n",
    "    executes it, and then returns the result. If any exception occurs during execution, it returns None.\n",
    "\n",
    "    Args:\n",
    "        input (str): The input to pass to 'function_call'.\n",
    "\n",
    "    Returns:\n",
    "        str: The result of the function call.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        func = function_call(input).strip()\n",
    "        print(func)\n",
    "        code = f\"result = {func}\"\n",
    "        local_vars = {}\n",
    "        exec(code, globals(), local_vars)\n",
    "        ans = local_vars.get(\"result\").strip()\n",
    "        return ans\n",
    "    except Exception as e:\n",
    "        return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = llm_pack(\"What the stock price for Tesla?\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = llm_pack(\"Who runs Tesla?\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = llm_pack(\"What does Tesla build?\")\n",
    "pprint.pprint(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = llm_pack(\"Who are the main competitors of Tesla?\")\n",
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
